#!/usr/bin/python3
# -*- coding:utf-8 -*-
import torch
import numpy as np
from einops import repeat

def gen_golden_data_simple():
    # Parse tiling data in ../main.cpp
    """
    const uint32_t aivNum = 8;
    const std::vector<uint32_t> shape = {32, 32};
    const std::string reduce_str = "add";
    const bool include_self = false;
    const uint32_t dim = 1;
    const uint32_t BLOCK_SIZE = 64;
    """
    block_num = 8
    self_shape = [32, 32]
    index_shape = [32, 32]
    reduce_str = "add"
    include_self = False
    dim = 1
    for line in open("./main.cpp", "r"):
        if "const uint32_t blockNum" in line:
            block_num = int(line.split()[-1].split(";")[0])
        if "const std::vector<uint32_t> selfShape" in line:
            self_shape = list(map(int, line.replace("{", "").replace("}", "").split(" = ")[-1].split(";")[0].split(",")))
        if "const std::vector<uint32_t> indexShape" in line:
            index_shape = list(map(int, line.replace("{", "").replace("}", "").split(" = ")[-1].split(";")[0].split(",")))
        if "const std::string reduce_str" in line:
            reduce_str = line.split()[-1].split(";")[0].replace("\"", "")
        if "const bool include_self" in line:
            include_self = line.split()[-1].split(";")[0] == "true"
        if "const uint32_t dim" in line:
            dim = int(line.split()[-1].split(";")[0])

    with open("./common.h", "r") as f:
        lines = f.readlines()
        # Parse typedef float DTYPE_SELF;
        for line in lines:
            if "typedef _Float32 DTYPE_SELF" in line:
                dtype = np.float32
                break
            elif "typedef _Float16 DTYPE_SELF" in line:
                dtype = np.float16
                break
            elif "typedef int32_t DTYPE_SELF" in line:
                dtype = np.int32
                break
            elif "typedef double DTYPE_SELF" in line:
                dtype = np.float64
                break
    
    print(f"block_num: {block_num}, self_shape: {self_shape}, index_shape: {index_shape}, reduce_str: {reduce_str}, include_self: {include_self}, dim: {dim}, dtype: {dtype}")
    '''
    input_self = np.random.uniform(1, 100, self_shape).astype(dtype)
    input_src = np.random.uniform(1,100, index_shape).astype(dtype)
    input_index = np.random.uniform(0, self_shape[dim], index_shape).astype(np.int32)

input_self: [19.52 52.3  29.34  4.85 71.06 96.3  35.38 44.44 91.8  87.94 22.92 91.56
 80.44 75.8  95.4  41.66 33.34 72.94]
input_index: [0 0 1 0 1 0 1 0 0 1 1 1 0 1 1 0 1 0 1 0 1 1 1 0 0 1 0 1 0 0 0 0 1 0 0 0 0
 1 1 0 1 1 1 0 0 0 0 0 1 1 1 0 0 1 0 0 1 0 1 1 1 0 0 0 1 1 0 0 1 0 0 1 1 0
 0 1 0 0 1 1 1 1 1 0 0 1 0 0 1 0 0 1 1 0 0 1 1 0 0 1 1 0 1 0 0 0 0 1]
input_src: [76.25  83.56  81.9   20.47  23.53  72.94  96.44  22.84   8.78  82.3
 57.1   83.94  85.06  97.25  39.34  42.5   61.28  54.66   2.586 10.375
 88.94  10.61  20.36  48.47  16.5   21.02   8.5    5.12  96.2   94.9
 31.75  59.12  66.25  83.4   74.7    4.83  24.55  83.44  51.66  71.06
 30.52  85.6   37.53  46.34  54.22  11.31   1.243  7.746 59.66  56.28
 82.56   1.571 64.7   84.2    5.734 17.95  64.06  19.4   87.9    3.367
 25.9   95.6   80.1   73.75  39.9    1.156 39.94   1.764 38.28  99.9
 30.89  46.5   16.81  64.3    4.816 96.    79.75  81.25   6.06  13.92
 79.4   65.06  54.3   39.3    7.133  8.3   38.84  51.12  20.39  81.94
 67.06  47.3    8.85  32.6   32.75   6.93  76.44  99.94  52.03  37.44
 32.2   83.06  80.56  79.    30.06  47.72  33.84  85.56 ]
Expected [43.47 68.2  45.1  40.66 54.1  47.2  31.48 55.53 33.53 65.75 55.62 40.53
 58.34 50.9  50.84 22.25 57.38 64.2 ], Shape (18,), Got [43.47 68.2  45.03 40.62 54.1  47.2  31.5  55.5  33.53 65.75 55.62 40.53
 58.28 50.9  50.84 22.23 57.34 64.2 ], Shape (18,)
Absolute error: [0.      0.      0.0625  0.03125 0.      0.      0.01563 0.03125 0.
 0.      0.      0.      0.0625  0.      0.      0.01563 0.03125 0.     ]
    '''
    # Reshape the arrays to match the dimensions in self_shape and index_shape
    input_self = np.array([19.52, 52.3, 29.34, 4.85, 71.06, 96.3, 35.38, 44.44, 91.8, 87.94, 22.92, 91.56,
                               80.44, 75.8, 95.4, 41.66, 33.34, 72.94], dtype=dtype).reshape(self_shape)
    input_index = np.array([0, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0,
 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 1, 0,
 0, 1, 0, 0, 1, 1, 1, 1, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1], dtype=np.int32).reshape(index_shape)
    input_src = np.array([76.25, 83.56, 81.9, 20.47, 23.53, 72.94, 96.44, 22.84, 8.78, 82.3,
 57.1, 83.94, 85.06, 97.25, 39.34, 42.5, 61.28, 54.66, 2.586, 10.375,
 88.94, 10.61, 20.36, 48.47, 16.5, 21.02, 8.5, 5.12, 96.2, 94.9,
 31.75, 59.12, 66.25, 83.4, 74.7, 4.83, 24.55, 83.44, 51.66, 71.06,
 30.52, 85.6, 37.53, 46.34, 54.22, 11.31, 1.243, 7.746, 59.66, 56.28,
 82.56, 1.571, 64.7, 84.2, 5.734, 17.95, 64.06, 19.4, 87.9, 3.367,
 25.9, 95.6, 80.1, 73.75, 39.9, 1.156, 39.94, 1.764, 38.28, 99.9,
 30.89, 46.5, 16.81, 64.3, 4.816, 96.0, 79.75, 81.25, 6.06, 13.92,
 79.4, 65.06, 54.3, 39.3, 7.133, 8.3, 38.84, 51.12, 20.39, 81.94,
 67.06, 47.3, 8.85, 32.6, 32.75, 6.93, 76.44, 99.94, 52.03, 37.44,
 32.2, 83.06, 80.56, 79.0, 30.06, 47.72, 33.84, 85.56], dtype=dtype).reshape(index_shape)
    
    input_self_cpu = torch.from_numpy(input_self)
    input_src_cpu = torch.from_numpy(input_src)
    input_index_cpu = torch.from_numpy(input_index).to(torch.long) # torch only supports int64
    
    cpu_result = torch.scatter_reduce(input=input_self_cpu,dim=dim,index=input_index_cpu,src=input_src_cpu,reduce=reduce_str,include_self=include_self)
    golden = cpu_result.numpy().astype(dtype)

    input_self.tofile("./input/input_self.bin")
    input_index.tofile("./input/input_index.bin")
    input_src.tofile("./input/input_src.bin")
    golden.tofile("./output/golden.bin")

if __name__ == "__main__":
    gen_golden_data_simple()
